# coding=utf-8
"""
author = jamon
"""
import asyncio
import struct
from asyncio import Future

import ujson

from obespoir.base.common_define import ConnectionStatus, RpcMsgType
from obespoir.base.global_object import GlobalObject
from obespoir.base.ob_exception import ObException
from obespoir.base.ob_protocol import ObProtocol
from obespoir.rpcserver.connection_manager import RpcConnectionManager
from obespoir.rpcserver.route import rpc_message_handle
from obespoir.rpcserver.rpc_index import RpcIndex
from obespoir.share.encodeutil import AesEncoder


class RpcPushProtocol(ObProtocol):

    def __init__(self):
        super().__init__()
        self.host = None
        self.port = None

        # (int, int, int)  -> (message_length, command_id, version)
        self.hand_fmt = "iii"
        self.head_len = struct.calcsize(self.hand_fmt)
        self.identifier = 0

        self.encode_ins = AesEncoder(
            GlobalObject().rpc_password,
            encode_type=GlobalObject().rpc_encode_type)
        self.version = 0

        self._buffer = b""    # 数据缓冲buffer

        # 消息头, list,   [message_length, command_id, version]
        self._head = None
        self.transport = None
        self.loop = asyncio.get_running_loop()

    def sync_default_keys(self):
        self.add_default_keys("addr", f"{self.host}:{self.port}")

    async def send_message(self, command_id, message, session_id,
                           to=None, msg_type=RpcMsgType.REQUEST):
        self.debug("start", message=message, session_id=session_id,
                   to=to, msg_type=msg_type)
        seq: str = RpcIndex.next()
        data = self.pack(message, command_id, session_id, to,
                         seq=seq, msg_type=msg_type)
        self.debug("end", message=message, session_id=session_id,
                   to=to, seq=seq, msg_type=msg_type)
        if msg_type == RpcMsgType.NOTIFY:
            self.transport.write(data)
            return
        wait_future: Future = self.loop.create_future()
        self.transport.write(data)
        RpcConnectionManager().add_future(seq, wait_future)
        return await wait_future

    async def message_handle(self, command_id, version, data):
        """
        实际处理消息
        :param command_id:
        :param version:
        :param data:
        :return:
        """
        self.debug("start", command_id=command_id, data=data)
        if not isinstance(data, dict):
            data = data.decode("utf-8") if isinstance(data, bytes) else data
            data = ujson.loads(data)
        session_id = data.get("session_id", None)  # 消息最初来源于哪个proxy节点
        seq = data.get("seq", None)
        if seq is None or seq not in RpcConnectionManager().future_dict:
            self.error("lose wait future", command_id=command_id, data=data)
            raise ObException(name="FutureNotExist", data=data)
        result = await rpc_message_handle(command_id, session_id, data)
        RpcConnectionManager().set_future_result_and_close(seq, result)
        self.debug("end", command_id=command_id, result=result)
        return

    def connection_made(self, transport):
        self.transport = transport
        address = transport.get_extra_info('peername')
        self.host, self.port = address
        RpcConnectionManager().store_connection(
            *address, self, status=ConnectionStatus.ESTABLISHED)
        self.debug('success', address=address)

    def data_received(self, data):
        self.debug('start', data=data)
        super().data_received(data)

    def eof_received(self):
        self.debug('start')
        if self.transport and self.transport.can_write_eof():
            self.transport.write_eof()
        RpcConnectionManager().lost_connection(self.host, self.port)

    def connnection_lost(self, exc):
        self.debug('start')
        RpcConnectionManager().lost_connection(self.host, self.port)
        super(RpcPushProtocol, self).connection_lost(exc)
